#!/usr/bin/env python3
"""
后处理Excel边表，过滤无效的invoke边
只保留满足以下条件的invoke边：
1. 同文件内的调用
2. target文件已被source文件导入的跨文件调用
"""

import pandas as pd
import os
from collections import defaultdict
import argparse


def extract_file_from_node(node_name):
    """从节点名称中提取文件路径"""
    if pd.isna(node_name):
        return None
    if ':' in str(node_name):
        return str(node_name).split(':')[0]
    return str(node_name)


def build_import_map(edges_df):
    """
    构建导入映射表
    返回: dict {source_file: set(imported_files)}
    """
    import_map = defaultdict(set)

    # 筛选import类型的边
    import_edges = edges_df[edges_df['type'] == 'imports']

    for _, row in import_edges.iterrows():
        source_file = extract_file_from_node(row['source'])
        target_file = extract_file_from_node(row['target'])

        if source_file and target_file:
            import_map[source_file].add(target_file)

    return import_map


def validate_invocation(source_node, target_node, import_map):
    """
    验证一个调用是否有效

    返回: (is_valid, reason)
    """
    source_file = extract_file_from_node(source_node)
    target_file = extract_file_from_node(target_node)

    if not source_file or not target_file:
        return False, "Invalid node format"

    # 情况1：同文件内的调用总是有效的
    if source_file == target_file:
        return True, "Same file invocation"

    # 情况2：检查target_file是否被source_file导入
    imported_files = import_map.get(source_file, set())
    if target_file in imported_files:
        return True, "Target file is imported"

    return False, f"File '{target_file}' not imported by '{source_file}'"


def post_process_edges(input_excel, output_dir):
    """
    后处理边表Excel文件
    """
    print(f"Loading edges from {input_excel}")
    edges_df = pd.read_excel(input_excel)

    print(f"Original edges: {len(edges_df)}")

    # 统计原始边类型
    original_stats = edges_df['type'].value_counts()
    print("\nOriginal edge statistics:")
    for edge_type, count in original_stats.items():
        print(f"  {edge_type}: {count:,}")

    # 构建导入映射
    print("\nBuilding import map...")
    import_map = build_import_map(edges_df)
    print(f"Found {len(import_map)} files with imports")

    # 显示一些导入示例
    print("\nSample import relationships:")
    for i, (source, targets) in enumerate(import_map.items()):
        if i >= 3:  # 只显示前3个
            break
        print(f"  {source} imports:")
        for target in list(targets)[:3]:  # 每个文件最多显示3个导入
            print(f"    - {target}")

    # 处理edges
    valid_edges = []
    invalid_invocations = []

    print("\nValidating invocation edges...")
    total_invokes = 0
    valid_invokes = 0
    invalid_invokes = 0

    for idx, row in edges_df.iterrows():
        if row['type'] == 'invokes':
            total_invokes += 1
            is_valid, reason = validate_invocation(row['source'], row['target'], import_map)

            if is_valid:
                valid_edges.append(row)
                valid_invokes += 1
            else:
                invalid_invokes += 1
                invalid_invocations.append({
                    'source': row['source'],
                    'target': row['target'],
                    'source_file': extract_file_from_node(row['source']),
                    'target_file': extract_file_from_node(row['target']),
                    'reason': reason
                })

                # 显示前几个无效调用作为示例
                if invalid_invokes <= 5:
                    print(f"  Invalid invoke #{invalid_invokes}:")
                    print(f"    {row['source']} -> {row['target']}")
                    print(f"    Reason: {reason}")
        else:
            # 其他类型的边直接保留
            valid_edges.append(row)

        # 显示进度
        if (idx + 1) % 10000 == 0:
            print(f"  Processed {idx + 1:,} / {len(edges_df):,} edges...")

    # 创建清理后的DataFrame
    cleaned_df = pd.DataFrame(valid_edges)

    # 统计清理后的边类型
    cleaned_stats = cleaned_df['type'].value_counts()

    print(f"\nProcessing complete!")
    print(f"  Total invocation edges: {total_invokes:,}")
    print(f"  Valid invocations: {valid_invokes:,}")
    print(f"  Invalid invocations removed: {invalid_invokes:,}")
    print(f"  Removal rate: {invalid_invokes/total_invokes*100:.1f}%")

    print("\nCleaned edge statistics:")
    for edge_type, count in cleaned_stats.items():
        original_count = original_stats.get(edge_type, 0)
        change = count - original_count
        change_pct = (change / original_count * 100) if original_count > 0 else 0
        print(f"  {edge_type}: {count:,} (change: {change:+,}, {change_pct:+.1f}%)")

    # 保存结果
    os.makedirs(output_dir, exist_ok=True)

    # 1. 保存清理后的边表
    output_edges = os.path.join(output_dir, 'edges_cleaned.xlsx')
    cleaned_df.to_excel(output_edges, index=False)
    print(f"\nCleaned edges saved to: {output_edges}")

    # 2. 保存无效调用详情
    if invalid_invocations:
        invalid_df = pd.DataFrame(invalid_invocations)

        # 添加统计信息
        invalid_stats = invalid_df.groupby('reason').size().reset_index(name='count')

        output_invalid = os.path.join(output_dir, 'invalid_invocations.xlsx')
        with pd.ExcelWriter(output_invalid) as writer:
            invalid_df.to_excel(writer, sheet_name='Details', index=False)
            invalid_stats.to_excel(writer, sheet_name='Statistics', index=False)

        print(f"Invalid invocations saved to: {output_invalid}")

    # 3. 保存统计对比
    statistics = {
        'Metric': [
            'Total Edges (Original)',
            'Total Edges (Cleaned)',
            'Edges Removed',
            '',
            'Invocation Edges (Original)',
            'Invocation Edges (Valid)',
            'Invocation Edges (Removed)',
            'Removal Rate',
            '',
            'Contains Edges',
            'Import Edges',
            'Inheritance Edges'
        ],
        'Count': [
            len(edges_df),
            len(cleaned_df),
            len(edges_df) - len(cleaned_df),
            '',
            total_invokes,
            valid_invokes,
            invalid_invokes,
            f"{invalid_invokes/total_invokes*100:.1f}%",
            '',
            cleaned_stats.get('contains', 0),
            cleaned_stats.get('imports', 0),
            cleaned_stats.get('inherits', 0)
        ]
    }

    stats_df = pd.DataFrame(statistics)
    output_stats = os.path.join(output_dir, 'processing_statistics.xlsx')
    stats_df.to_excel(output_stats, index=False)
    print(f"Statistics saved to: {output_stats}")

    # 显示最终汇总
    print("\n" + "="*60)
    print("SUMMARY")
    print("="*60)
    print(f"Original edges: {len(edges_df):,}")
    print(f"Cleaned edges: {len(cleaned_df):,}")
    print(f"Edges removed: {len(edges_df) - len(cleaned_df):,}")
    print(f"Output directory: {output_dir}")

    return cleaned_df, invalid_invocations


def main():
    parser = argparse.ArgumentParser(
        description='Post-process Excel edge table to remove invalid invocations'
    )
    parser.add_argument(
        '--input',
        type=str,
        required=True,
        help='Path to input Excel file (edges_info.xlsx)'
    )
    parser.add_argument(
        '--output-dir',
        type=str,
        required=True,
        help='Output directory for cleaned data'
    )

    args = parser.parse_args()

    if not os.path.exists(args.input):
        print(f"Error: Input file not found: {args.input}")
        return 1

    post_process_edges(args.input, args.output_dir)
    return 0


if __name__ == "__main__":
    import sys

    # 如果没有提供参数，使用默认值
    if len(sys.argv) == 1:
        input_file = "/data_ext/ref_code/LocAgent/index_data/pp_data/graph_index_v2.3/pingp-auto-test_analysis/edges_info.xlsx"
        output_dir = "/data_ext/ref_code/LocAgent/index_data_cleaned/pingp-auto-test"

        if os.path.exists(input_file):
            print(f"Processing default file: {input_file}")
            print(f"Output directory: {output_dir}")
            print("-" * 60)
            post_process_edges(input_file, output_dir)
        else:
            print(f"Default file not found: {input_file}")
            print("Please provide input file path using --input parameter")
            sys.exit(1)
    else:
        sys.exit(main())